#!/usr/bin/env python
#
# testing for keep-alive connection
#   agent-pool support
#

import time,socket,base64,fcntl,errno
import threading
import Queue as queue

import celtuce.net.ioloop as ioloop
import celtuce.net.iostream as iostream

from celtuce.help import logging
logging.getLogger().setLevel(logging.LEVEL_DEBUG)


URL = 'tcp://127.0.0.1:14300'

#######################################################################


class   ServerThread(threading.Thread) :
    def __init__(self, handle_callable, tcp_port=0) :
        threading.Thread.__init__(self) 

        self._tcp_server = None

        self._io_loop = ioloop.IOLoop()

        self._agent_pool = AgentPool(self._io_loop)

        if tcp_port :
            self._tcp_port = tcp_port
            self._tcp_server = TcpServer(self._agent_pool.handleClientRequest, self._io_loop)

        self._logger = logging.getLogger('serverthread')
        ## __init__()

    def run(self) :
        self._tcp_server.listen( self._tcp_port )

        self._io_loop.start()
        ## run()

    def stop(self) :
        self._logger.info('stop IOLoop ...')
        self._io_loop.stop()
        ## stop()

    ## class ServerThread


class   TcpServer(object) :
    def __init__(self, request_callback, io_loop) :
        self._request_callback = request_callback
        self._io_loop = io_loop

        self._socket = None
        self._connections = {}
        self._logger = logging.getLogger('tcpserver')
        ## __init__()

    def listen(self, port, address='') :
        assert not self._socket
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        flags = fcntl.fcntl(self._socket.fileno(), fcntl.F_GETFD)
        flags |= fcntl.FD_CLOEXEC
        fcntl.fcntl(self._socket.fileno(), fcntl.F_SETFD, flags)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._socket.setblocking(0)
        self._socket.bind((address, port))
        self._socket.listen(128)
        self._io_loop.add_handler(self._socket.fileno(), self._handleEvents, self._io_loop.READ)
        ## listen()

    def _handleEvents(self, fd, events):
        while True:
            try:
                connection, address = self._socket.accept()
            except socket.error, e:
                if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                    return
                raise
            try:
                stream = iostream.IOStream(connection, io_loop=self._io_loop)
                self._connections[ address ] = TcpConnection(stream, address, self._request_callback)
            except:
                self._logger.excpt()
        ## _handleEvents()


    ## class TcpServer

class   TcpConnection(object) :
    def __init__(self, stream, address, request_callback) :
        self._stream = stream
        self._address = address
        self._request_callback = request_callback

        self._END_TAG = "\r\n\r\n"
        self._request_finished = False
        self._stream.read_until(self._END_TAG, self.__onRequest)

        self._keep_alive = True

        self._logger = logging.getLogger('tcpconnect')
        self._logger.debug('create TcpConnection for %s', self._address)
        ## __init__()

    def __del__(self) :
        #if self._stream.closed() :
        #    self._logger.debug('stream is closed when destroying')
        #else :
        #    self._logger.debug('stream is opened when destroying')
        self._logger.debug('destroy TcpConnection for %s', self._address)
        pass
        ## __del__()

    def __onRequest(self, data) :
        request = data
        #self._logger.debug('recv request  : %s', repr(request))
        #response = self._request_callback(request)
        self._request_callback(self, request)
        ## __onRequest()

    def sendResponse(self, response) :
        self.write(response)
        ## sendResponse()

    def write(self, data) :
        if data :
            response = data 
            #self._logger.debug('send response : %s', repr(response))
            self._stream.write(''.join((response, self._END_TAG)), self.__onWriteComplete)
        else :
            if not self._keep_alive :
                self.finish()
        ## write()

    def finish(self):
        if self._stream and not self._stream.writing():
            self.__finishRequest()

    def __onWriteComplete(self):
        if self._keep_alive :
            self._request_finished = False
            self._stream.read_until(self._END_TAG, self.__onRequest)
        else :
            self.__finishRequest()
        ## __onWriteComplete()

    def __finishRequest(self) :
        #self._logger.debug('finish request ...')
        self._stream.close()
        self._stream = None
        ## __finishRequest()

    ## class TcpConnection


#######################################################################

try :
    from hashlib import md5 as MD5
except :
    from md5 import md5 as MD5

class   AgentPool(object) :

    def __init__(self, io_loop) :
        self._io_loop = io_loop

        self._agents = {}
        ## __init__()

    def handleClientRequest(self, client_connection, request) :
        # STEP. parse packet, to get OP,Key,timestamp,etc ...
        # TODO

        # STEP. lookup agent for this request
        agent = self.lookupAgent( request )

        # STEP. put request into queue of the agent
        agent.sendRequest( client_connection, request )

        ## handleClientRequest()

    def lookupAgent(self, request) :
        id = MD5(request).hexdigest()[:1]

        if id not in self._agents :
            q = queue.Queue()
            agent = Agent(q)
            self._agents[ id ] = agent

        return self._agents[ id ]
        ## request()

    ## class AgentPool

class   Agent(object) :
    def __init__(self, incoming_queue) :
        self._q = incoming_queue

        self._client_connection = None
        ## __init__()

    def sendRequest(self, client_connection, request) :
        if self._q.empty() :
            response = request
            client_connection.sendResponse(response)
            return

        self._q.put( request )
        while True :
            try :
                response = q.get(False) 
                client_connection.sendResponse(response)
                self._q.task_done()
            except queue.Empty :
                break
        ## sendRequest()

    ## class Agent

#######################################################################

class   Packet(object)  :

    MIN_PACKET_ID = 1
    MAX_PACKET_ID = 65535
    packet_id_current = MAX_PACKET_ID

    ACTION_REQUEST  = 'REQUEST'
    ACTION_RESPONSE = 'RESPONSE'

    OP_APPEND = 'APPEND'
    OP_INSERT = 'INSERT'
    OP_UPDATE = 'UPDATE'
    OP_REPLACE = 'REPLACE'
    OP_ERASE  = 'ERASE'
    OP_FIND   = 'FIND'

    def __init__(self) :
        self.clear()
        ## __init__()

    def clear(self) :
        self.action = None
        self.table = None
        self.op = None
        self.key = None
        self.data = None
        self.error = None
        self.state = 0
        self.packid = 0
        self.respid = 0
        ## clear()

    def __str__(self) :
        s = "action:%s table:%s op:%s key:%s packid:%d" % (self.action, self.table, self.op, self.key, self.packid)
        if self.respid > 0 :
            s += " respid:%d" % self.respid
        if self.state :
            s += ' state:%d' % self.state

        if self.data :
            s += ' data:%s' % self.data
        if self.error :
            s += ' error:%s' % self.error
        return s
        ## __str__()

    def isComplete(self) :
        if not self.action or not self.table or not self.op or not self.key :
            return  False
        return  True

    def setRequest(self) :
        if self.ACTION_REQUEST == self.action : return
        self.action = self.ACTION_REQUEST
    def setResponse(self) :
        if self.ACTION_RESPONSE == self.action : return
        self.action = self.ACTION_RESPONSE

    def isRequest(self) :
        if self.ACTION_REQUEST == self.action : return True
        return  False
    def isResponse(self) :
        if self.ACTION_RESPONSE == self.action : return True
        return  False

    ####-----------------------------------------------------------####

    def depack(cls, record) :
        '''decode packet record string into Packet struct'''
        field_separator = '\t'
        kv_separator = ':'
        packet = Packet()
        fields = record.split(field_separator)
        for f in fields :
            (key, value) = f.split(kv_separator, 1)
            if not value : continue
            if 'action' == key :
                packet.action = value
            elif 'table' == key :
                packet.table = value
            elif 'op' == key :
                packet.op = value
            elif 'key' == key :
                packet.key = base64.b64decode(value)
            elif 'data' == key :
                packet.data = base64.b64decode(value)
            elif 'error' == key :
                packet.error = base64.b64decode(value)
            elif 'state' == key :
                packet.state = int(value)
            elif 'packid' == key :
                packet.packid = int(value)
            elif 'respid' == key :
                packet.respid = int(value)

        if not packet.isComplete() : return None
        return  packet
        # depack()
    depack = classmethod(depack)

    def enpack(cls, packet) :
        '''encode Packet struct into packet record string'''
        packid = packet.packid
        if not packid :  packid = cls.nextPacketID()
        fields = []
        fields.append('action:%s' % packet.action)
        fields.append('table:%s' % packet.table)
        fields.append('packid:%d' % packid)
        if packet.respid > 0 :
            fields.append('respid:%d' % packet.respid)
        fields.append('op:%s' % packet.op)
        fields.append('state:%d' % packet.state)
        fields.append(':'.join(['key', base64.b64encode(packet.key)]))
        if packet.data :
            fields.append(':'.join(['data', base64.b64encode(packet.data)]))
        if packet.error :
            fields.append(':'.join(['error', base64.b64encode(packet.error)]))
        record = '\t'.join(fields)
        return  (record, packid)
        ## enpack()
    enpack = classmethod(enpack)

    def nextPacketID(cls) :
        if cls.packet_id_current == cls.MAX_PACKET_ID :
            cls.packet_id_current = cls.MIN_PACKET_ID
        else :
            cls.packet_id_current += 1
        return  cls.packet_id_current
        ## nextPacketID()
    nextPacketID = classmethod(nextPacketID)

    ## class Packet


class   PacketTuple(object)  :

    MIN_PACKET_ID = 1
    MAX_PACKET_ID = 65535
    packet_id_current = MAX_PACKET_ID

    ACTION_REQUEST  = 'REQUEST'
    ACTION_RESPONSE = 'RESPONSE'

    OP_APPEND = 'APPEND'
    OP_INSERT = 'INSERT'
    OP_UPDATE = 'UPDATE'
    OP_REPLACE = 'REPLACE'
    OP_ERASE  = 'ERASE'
    OP_FIND   = 'FIND'

    def __init__(self) :
        self.clear()
        ## __init__()

    def clear(self) :
        #            action table op    key   data  error state packid respid
        self._list = [None, None, None, None, '',   '',    0,    0,     0]
        ## clear()

    def setAction(self, action) :
        self._list[0] = action
    def getAction(self) :
        return self._list[0]
    action = property(getAction, setAction)

    def setTable(self, table) :
        self._list[1] = table
    def getTable(self) :
        return self._list[1]
    table = property(getTable, setTable)

    def setOp(self, op) : 
        self._list[2] = op
    def getOp(self) :
        return self._list[2]
    op = property(getOp, setOp)

    def setKey(self, key) :
        self._list[3] = key
    def getKey(self) :
        return self._list[3]
    key = property(getKey, setKey)

    def setData(self, data) :
        self._list[4] = data
    def getData(self) :
        return self._list[4]
    data = property(getData, setData)

    def setError(self, error) :
        self._list[5] = error
    def getError(self) :
        return self._list[5]

    def setState(self, state) :
        self._list[6] = state
    def getState(self) :
        return self._list[6]

    def enpack(cls, packet) :
        to_pack = [ str(x) for x in packet._list ]
        to_pack[ 3 ] = base64.b64encode( to_pack[3] )
        to_pack[ 4 ] = base64.b64encode( to_pack[4] )
        to_pack[ 5 ] = base64.b64encode( to_pack[5] )

        record = '\t'.join(to_pack)
        return (record,0)
        ## enpack()
    enpack = classmethod(enpack)

    def depack(cls, record) :
        pack = []
        pack.extend( record.split('\t') )

        pack[ 3 ] = base64.b64decode( pack[3] )
        pack[ 4 ] = base64.b64decode( pack[4] )
        pack[ 5 ] = base64.b64decode( pack[5] )
        pack[ 6 ] = int( pack[6] )
        pack[ 7 ] = int( pack[7] )
        pack[ 8 ] = int( pack[8] )
        return pack
        ## depack()
    depack = classmethod(depack)

    ## class Packet
#######################################################################

class   Client(object) :
    def __init__(self, server_host, server_port) :
        self._server_host = server_host
        self._server_port = server_port

        self._END_TAG = "\r\n\r\n"

        self._socket = None

        self._logger = logging.getLogger('client')
        ## __init__()

    def __del__(self) :
        self.__disconnect()
        ## __del__()

    ####-----------------------------------------------------------####

    def __connect(self) :
        if self._socket : 
            #self._stream = iostream.IOStream( self._socket, self._io_loop )
            return

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._server_host, self._server_port) )
        except :
            self._logger.excpt()
            raise
        self._socket = s

        self._io_loop = ioloop.IOLoop()
        self._stream = iostream.IOStream( self._socket, self._io_loop )

        return
        ## __connect()

    def __disconnect(self) :
        if self._stream :
            self._stream.close()
            self._stream = None

        if self._socket :
            try :
                self._socket.close()
            except :
                self._logger.excpt()

        self._socket = None
        self._io_loop = None
        ## __disconnect()

    def rpc(self, request) :
        self.__connect()

        #self._logger.debug('write request : %s', request)
        self._stream.write( ''.join((request, self._END_TAG)), self.__onWriteComplete )
        self._stream.read_until(self._END_TAG, self.__onResponse)
        
        self._io_loop.start()

        return self._response
        ## rpc()

    def __onWriteComplete(self) :
        #self._logger.debug('write done.')
        pass
        ## __onWriteComplete()

    def __onResponse(self, response_str) :
        #self._logger.debug('read response : %s', repr(response_str))
        response = response_str.strip()
        if response :
            self._response = response_str.strip()
            self._io_loop.stop()
        else :
            self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __onResponse()

    ## class Client


#######################################################################


def echo(any_data) :
    return any_data.strip()
    ## echo()


TCP_PORT = 24344

def runTest() :
    logger = logging.getLogger('test')

    logger.info('test keep-alive connection startup ...')
    server = ServerThread(echo, TCP_PORT)
    server.start()

    time.sleep(0.5)

    client = Client('', TCP_PORT)

    try :
        # ping
        request = 'hello,world'
        response = client.rpc( request )
        logger.info('REQUEST  : %s', request)
        logger.info('RESPONSE : %s', response)

        # test tps
        key_prefix = 'hello'
        data_prefix = 'world'
        time_start = time.time()
        times = 100000
        sum_bytes = 0
        for i in range(times) :
            packet = Packet()
            packet.action = Packet.ACTION_REQUEST
            packet.table = 'test.gate4'
            packet.op = Packet.OP_INSERT
            packet.key = '-'.join((key_prefix, str(i)))
            packet.data = '-'.join((data_prefix, str(i)))
            request,request_id = Packet.enpack( packet )
            response_str = client.rpc( request )
            response = Packet.depack( response_str )

            sum_bytes += len(response_str)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s',times,time_last,tps,bw)
    finally :
        logger.info('shutdown ...')
        server.stop()
        logger.info('joining server-thread ...')
        if server.isAlive() :
            server.join()

    logger.info('test keep-alive connection done')
    ## runTest()

if  __name__ == '__main__' :
    runTest()

